#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
本地人脸识别服务（LBPH + Haar 检测 + 简易人脸库）
 - 依赖: OpenCV (建议安装 opencv-contrib-python 以启用 cv2.face)
 - 存储: faces/ 目录下的人脸图片与模型，labels.json 记录标签与姓名映射
"""

import os
import json
from typing import Dict, List, Optional, Tuple

import cv2
import numpy as np


class FaceRecognizer:
    """本地人脸识别服务，支持新增样本、训练与识别。"""

    def __init__(self, base_dir: str = "faces") -> None:
        self.base_dir = base_dir
        self.labels_path = os.path.join(self.base_dir, "labels.json")
        self.model_path = os.path.join(self.base_dir, "lbph_model.yml")
        self.haar_path = cv2.data.haarcascades + "haarcascade_frontalface_default.xml"
        self.face_cascade = cv2.CascadeClassifier(self.haar_path)

        os.makedirs(self.base_dir, exist_ok=True)
        self.labels: Dict[str, int] = self._load_labels()
        self.inverse_labels: Dict[int, str] = {v: k for k, v in self.labels.items()}

        self.recognizer = None
        self._init_recognizer()
        self._try_load_model()

        # 训练免模式：pHash 直匹配索引（列表项为 (phash:int, name:str)）
        self._phash_entries: List[Tuple[int, str]] = []
        try:
            self._build_phash_index()
        except Exception:
            self._phash_entries = []

    def _init_recognizer(self) -> None:
        try:
            self.recognizer = cv2.face.LBPHFaceRecognizer_create(radius=1, neighbors=8, grid_x=8, grid_y=8)
        except Exception:
            self.recognizer = None

    def _load_labels(self) -> Dict[str, int]:
        if os.path.exists(self.labels_path):
            try:
                with open(self.labels_path, "r", encoding="utf-8") as f:
                    data = json.load(f)
                    if isinstance(data, dict):
                        return {str(k): int(v) for k, v in data.items()}
            except Exception:
                pass
        return {}

    def _save_labels(self) -> None:
        try:
            with open(self.labels_path, "w", encoding="utf-8") as f:
                json.dump(self.labels, f, ensure_ascii=False, indent=2)
        except Exception:
            pass

    def _try_load_model(self) -> None:
        if self.recognizer is not None and os.path.exists(self.model_path):
            try:
                self.recognizer.read(self.model_path)
            except Exception:
                pass

    def _ensure_person_dir(self, name: str) -> str:
        safe = name.strip().replace("/", "_").replace("\\", "_")
        person_dir = os.path.join(self.base_dir, safe)
        os.makedirs(person_dir, exist_ok=True)
        return person_dir

    def add_face_image(self, name: str, bgr_image: np.ndarray) -> Optional[str]:
        """将上传的人脸图保存至库中（自动检测并裁剪），返回保存路径。"""
        if bgr_image is None or bgr_image.size == 0:
            return None

        gray = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2GRAY)
        # 增强对比度以提升检测与识别鲁棒性
        try:
            clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
            gray = clahe.apply(gray)
        except Exception:
            pass
        faces = self.face_cascade.detectMultiScale(gray, scaleFactor=1.2, minNeighbors=5, minSize=(60, 60))
        if len(faces) == 0:
            # 未检测到人脸，启用兜底：取图像中心区域作为近似人脸，提高样本入库成功率
            h, w = gray.shape[:2]
            side = min(h, w)
            cx, cy = w // 2, h // 2
            half = max(60, int(0.35 * side))
            x0 = max(0, cx - half)
            y0 = max(0, cy - half)
            x1 = min(w, cx + half)
            y1 = min(h, cy + half)
            crop = gray[y0:y1, x0:x1]
        else:
            # 取最大的人脸并适度扩展
            x, y, w, h = max(faces, key=lambda r: r[2] * r[3])
            pad = int(0.15 * max(w, h))
            x0 = max(0, x - pad)
            y0 = max(0, y - pad)
            x1 = min(gray.shape[1], x + w + pad)
            y1 = min(gray.shape[0], y + h + pad)
            crop = gray[y0:y1, x0:x1]

        person_dir = self._ensure_person_dir(name)
        idx = 1
        while True:
            candidate = os.path.join(person_dir, f"sample_{idx:03d}.png")
            if not os.path.exists(candidate):
                break
            idx += 1
        try:
            # 统一尺寸
            face_img = cv2.resize(crop, (200, 200))
            cv2.imwrite(candidate, face_img)
            # 确保label存在
            if name not in self.labels:
                new_id = (max(self.labels.values()) + 1) if self.labels else 1
                self.labels[name] = new_id
                self.inverse_labels[new_id] = name
                self._save_labels()
            # 更新直匹配索引
            try:
                ph = self._compute_phash(face_img)
                self._phash_entries.append((ph, name))
            except Exception:
                pass
            return candidate
        except Exception:
            return None

    def _collect_training_data(self) -> Tuple[List[np.ndarray], List[int]]:
        images: List[np.ndarray] = []
        label_ids: List[int] = []
        for name, lid in self.labels.items():
            person_dir = self._ensure_person_dir(name)
            for fn in os.listdir(person_dir):
                if not fn.lower().endswith((".png", ".jpg", ".jpeg")):
                    continue
                path = os.path.join(person_dir, fn)
                try:
                    img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
                    if img is None:
                        continue
                    # 若图片不是 200x200 做一次 resize
                    if img.shape[:2] != (200, 200):
                        img = cv2.resize(img, (200, 200))
                    images.append(img)
                    label_ids.append(lid)
                except Exception:
                    continue
        return images, label_ids

    def train(self) -> bool:
        """训练并保存模型。返回是否成功。"""
        if self.recognizer is None:
            return False
        images, label_ids = self._collect_training_data()
        if not images or not label_ids:
            return False
        # 放宽要求：至少 1 人、每人 ≥ 1 张（建议≥10张更稳定）
        person_counts: Dict[int, int] = {}
        for lid in label_ids:
            person_counts[lid] = person_counts.get(lid, 0) + 1
        if len(person_counts) < 1 or any(c < 1 for c in person_counts.values()):
            return False
        try:
            self.recognizer.train(images, np.array(label_ids))
            os.makedirs(self.base_dir, exist_ok=True)
            self.recognizer.write(self.model_path)
            return True
        except Exception:
            return False

    def recognize_in_bgr(self, bgr_image: np.ndarray) -> Optional[Tuple[str, float]]:
        """在整幅BGR图像中寻找人脸并识别，返回 (name, score) 或 None。
        优先使用免训练直匹配（pHash），若失败且LBPH可用则回退LBPH。
        约定: 直匹配返回的 score 为 Hamming 距离（越小越好）；LBPH 同样越小越好。
        """
        if bgr_image is None or bgr_image.size == 0:
            return None
        gray = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2GRAY)
        try:
            clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
            gray = clahe.apply(gray)
        except Exception:
            pass
        faces = self.face_cascade.detectMultiScale(gray, scaleFactor=1.2, minNeighbors=5, minSize=(60, 60))
        if len(faces) == 0:
            return None
        # 取最大脸
        x, y, w, h = max(faces, key=lambda r: r[2] * r[3])
        crop = gray[y:y+h, x:x+w]
        try:
            crop = cv2.resize(crop, (200, 200))
        except Exception:
            return None
        # 1) 直匹配 pHash
        try:
            ph = self._compute_phash(crop)
            best_name, best_dist = self._phash_search(ph)
            if best_name is not None and best_dist is not None:
                # 阈值：<= 12 认为较相似，<= 16 勉强接受
                if best_dist <= 16:
                    return best_name, float(best_dist)
        except Exception:
            pass

        # 2) 回退 LBPH（若可用）
        try:
            if self.recognizer is None:
                return None
            label_id, confidence = self.recognizer.predict(crop)
            name = self.inverse_labels.get(label_id)
            if name is None:
                return None
            if confidence < 120.0:
                return name, float(confidence)
            return None
        except Exception:
            return None

    def recognize_face_from_roi(self, roi_bgr: np.ndarray) -> Optional[Tuple[str, float]]:
        """对给定ROI识别人脸，返回 (name, confidence) 或 None。"""
        return self.recognize_in_bgr(roi_bgr)


    # -------------------- pHash 直匹配支持 --------------------
    def _compute_phash(self, gray_img: np.ndarray) -> int:
        """计算图像 pHash（返回64位整数）。输入需为灰度图 200x200。"""
        try:
            img = cv2.resize(gray_img, (32, 32))
        except Exception:
            img = gray_img
        img = np.float32(img)
        try:
            dct = cv2.dct(img)
        except Exception:
            # 退化：使用numpy实现DCT近似
            dct = np.fft.fft2(img)
            dct = np.abs(dct)
        # 取左上角 8x8 的低频成分
        dct_low = dct[:8, :8]
        # 均值（排除 DC 分量[0,0]）
        dct_flat = dct_low.flatten()
        avg = (np.sum(dct_flat) - dct_low[0, 0]) / (dct_flat.size - 1)
        bits = (dct_low >= avg).astype(np.uint8).flatten()
        # 转64位整数
        phash = 0
        for b in bits:
            phash = (phash << 1) | int(b)
        return int(phash)

    @staticmethod
    def _hamming(a: int, b: int) -> int:
        return int(bin(a ^ b).count("1"))

    def _build_phash_index(self) -> None:
        """扫描人脸库，构建 pHash 索引。"""
        entries: List[Tuple[int, str]] = []
        for name in self.labels.keys():
            person_dir = self._ensure_person_dir(name)
            for fn in os.listdir(person_dir):
                if not fn.lower().endswith((".png", ".jpg", ".jpeg")):
                    continue
                path = os.path.join(person_dir, fn)
                try:
                    img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
                    if img is None:
                        continue
                    if img.shape[:2] != (200, 200):
                        img = cv2.resize(img, (200, 200))
                    ph = self._compute_phash(img)
                    entries.append((ph, name))
                except Exception:
                    continue
        self._phash_entries = entries

    def _phash_search(self, ph: int) -> Tuple[Optional[str], Optional[int]]:
        """在索引中查找最小汉明距离的样本。返回 (name, dist)。"""
        if not self._phash_entries:
            try:
                self._build_phash_index()
            except Exception:
                return None, None
        best_name: Optional[str] = None
        best_dist: Optional[int] = None
        for eph, ename in self._phash_entries:
            d = self._hamming(ph, eph)
            if best_dist is None or d < best_dist:
                best_dist = d
                best_name = ename
        return best_name, best_dist

    def rebuild_index(self) -> None:
        """供外部调用，重新构建 pHash 索引（如删除/新增样本后）。"""
        try:
            self._build_phash_index()
        except Exception:
            self._phash_entries = []


